#!/usr/bin/python
'''Probability distribution function.'''

import random

def fermi_dirac(sequences, total_samples):
    '''Returns list of sequences based on based on Fermi-Dirac distribution.
    sequences is a list of sequences out of which we need to generate
    samples. Percent is percent sequences out of given sequences as sample.'''
    
    samples = []
    
    while len(samples) < total_samples:
        sample = random.choice(sequences)
        samples.append(sample)
        
    return samples

def boltzmann(energy_matrix, sequences, total_samples):
    '''Returns list of sequences based on based on Boltzmann distribution.
    matrix is matrix of real base frequency at various positions from which
    sample is to be generated. Actually it is just linear list here, but used as
    an array. sequences is a list of sequences out of which we need to generate
    samples. Percent is percent sequences out of given sequences as sample.'''

    for i in xrange(len(energy_matrix)):
        for j in xrange(len(energy_matrix[0])):
            energy_matrix[i][j] = pow(2, energy_matrix[i][j])

    for j in xrange(len(energy_matrix[0])):
        total = 0
        for i in xrange(len(energy_matrix)):
            total += energy_matrix[i][j]
        for i in xrange(len(energy_matrix)):
            energy_matrix[i][j] /= total

##     for i in xrange(len(energy_matrix)):
##         for j in xrange(len(energy_matrix[0])):
##             print "%.2f" % energy_matrix[i][j],

##         print

    freq_matrix = []

    for seq in energy_matrix:
        freq_matrix += seq

    bases = {'a':0, 'c':1, 'g':2, 't':3, 'A':0, 'C':1, 'G':2, 'T':3}
    samples = []
    samples_prob = {}
    total_prob = 0.0
    prob = 1.0
    
    for seq in sequences:
        i = 0
        prob = 1.0
        for base in seq:
            prob *= freq_matrix[7*bases[base]+i]
            i += 1
        total_prob += prob
        samples_prob[seq] = prob
        
    while total_samples > 0:
        rand_num = random.random() * total_prob
        temp = 0.0
        
        for seq in samples_prob.keys():
            temp += samples_prob[seq]
            if temp > rand_num:
                samples.append(seq)
                total_samples -= 1
                break
        
    return samples
        
def main():
    pass
    
if __name__ == "__main__":
    main()
    
